import numpy as np
import gradio as gr
import joblib
from skimage.transform import resize
import matplotlib.pyplot as plt
import os

# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# ---------------------- 1. 加载模型（增加有效性校验）----------------------
model = None
try:
    model = joblib.load('best_knn_model.pkl')
    # 校验模型是否可用（检查是否有predict方法）
    if hasattr(model, 'predict'):
        print("✅ 模型加载成功（可用）")
    else:
        raise ValueError("模型文件损坏，无predict方法")
except FileNotFoundError:
    print("❌ 未找到模型文件 'best_knn_model.pkl'，请先运行 optimal_knn.py 生成模型")
except Exception as e:
    print(f"❌ 模型加载错误: {str(e)}")

# ---------------------- 2. 图像预处理（与训练数据严格对齐）----------------------
def preprocess_image(image):
    """
    预处理逻辑：将画布图像转为MNIST格式（28x28灰度图、0-1归一化、黑底白字）
    返回：预处理后的一维数组 / 错误信息
    """
    try:
        # 步骤1：检查图像是否存在
        if image is None:
            return None, "请在画布上绘制一个数字（0-9）"
        
        # 步骤2：处理画布返回格式（兼容Gradio新旧版本）
        # 若为字典（新版本），提取image字段；若为数组（旧版本），直接使用
        if isinstance(image, dict):
            image = image.get("image", None)
            if image is None:
                return None, "未从画布提取到图像数据，请重新绘制"
        
        # 步骤3：转为灰度图（处理RGBA/RGB格式）
        if image.shape[-1] == 4:  # RGBA（含透明通道）
            # 提取RGB通道，按亮度公式转为灰度（更符合人眼感知）
            image_gray = np.dot(image[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
        elif len(image.shape) == 3 and image.shape[2] == 3:  # RGB
            image_gray = np.dot(image, [0.2989, 0.5870, 0.1140]).astype(np.uint8)
        else:  # 已为灰度图
            image_gray = image.astype(np.uint8)
        
        # 步骤4：检查是否有有效绘制（排除全白画布）
        if np.all(image_gray >= 250):  # 允许轻微灰底，避免误判
            return None, "未检测到有效绘制，请用深色笔在画布中央写清晰数字"
        
        # 步骤5：颜色反转（画布是“白底黑字”→ 转为MNIST“黑底白字”）
        image_inverted = 255 - image_gray  # 数字变白色（255），背景变黑色（0）
        
        # 步骤6：缩放为28x28（严格匹配MNIST尺寸）
        image_resized = resize(
            image_inverted, 
            (28, 28), 
            anti_aliasing=True,  # 抗锯齿，避免边缘模糊
            preserve_range=True  # 保留0-255像素范围
        ).astype(np.uint8)
        
        # 步骤7：归一化到0-1（与训练数据一致）
        image_scaled = image_resized / 255.0
        
        # 步骤8：展平为784维数组（模型输入格式）
        image_flattened = image_scaled.flatten()
        
        # 步骤9：保存预处理图像（供调试，可选）
        save_dir = "processed_images"
        os.makedirs(save_dir, exist_ok=True)
        plt.imsave(
            os.path.join(save_dir, "latest_processed.png"), 
            image_resized, 
            cmap='gray', 
            dpi=100
        )
        
        return image_flattened, None  # 成功：返回处理后的数据
    
    except Exception as e:
        error_msg = f"图像处理失败：{str(e)}"
        print(f"❌ {error_msg}")
        return None, error_msg

# ---------------------- 3. 预测函数（增加模型校验和置信度过滤）----------------------
def predict_digit(sketchpad_output):
    """预测手写数字，含模型校验和用户提示"""
    # 先检查模型是否可用
    if model is None:
        return "⚠️ 模型未加载成功，请先运行 optimal_knn.py 生成 'best_knn_model.pkl' 文件"
    
    # 预处理图像
    processed_image, error_msg = preprocess_image(sketchpad_output)
    if error_msg:
        return f"⚠️ {error_msg}"
    if processed_image is None:
        return "⚠️ 图像处理失败，请重新绘制数字"
    
    try:
        # 调整输入维度（模型需要(batch_size, 784)）
        input_data = processed_image.reshape(1, -1)
        
        # 预测数字
        prediction = model.predict(input_data)[0]
        
        # 计算置信度（KNN用“邻居平均距离”衡量，距离越小置信度越高）
        if hasattr(model, "kneighbors"):
            distances, _ = model.kneighbors(input_data)
            mean_distance = np.mean(distances)
            # 置信度映射到0-1（距离<5→置信度>0.8，距离>20→置信度<0.3）
            confidence = 1.0 / (1.0 + mean_distance / 5)  # 调整系数使置信度更合理
            confidence = max(0.01, min(0.99, confidence))  # 限制在0.01-0.99，避免极端值
        
            # 低置信度提示
            if confidence < 0.3:
                return f"预测数字：{int(prediction)}\n置信度：{confidence:.2f}\n⚠️ 置信度较低，请尽量写清晰数字（居中、笔画粗一些）"
            else:
                return f"预测数字：{int(prediction)}\n置信度：{confidence:.2f}\n✅ 识别成功"
        else:
            return f"预测数字：{int(prediction)}\n✅ 识别成功（无置信度计算）"
    
    except Exception as e:
        error_msg = f"预测失败：{str(e)}"
        print(f"❌ {error_msg}")
        return f"⚠️ {error_msg}"

# ---------------------- 4. Gradio界面（优化用户体验）----------------------
with gr.Blocks(title="手写数字识别器（MNIST-KNN）") as iface:
    gr.Markdown("# 🖌️ 手写数字识别器（0-9）")
    gr.Markdown("基于MNIST数据集训练的KNN模型，支持实时识别手写数字")
    
    with gr.Row():
        # 左侧：画布和操作按钮
        with gr.Column(scale=1):
            canvas = gr.Sketchpad(
                height=300,
                width=300,
                label="请在此处绘制数字"                
            )
            # 按钮组（横向排列）
            with gr.Row():
                clear_btn = gr.Button("清除画布", variant="secondary")
                predict_btn = gr.Button("开始识别", variant="primary")
        
        # 右侧：结果显示
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="识别结果",
                lines=4,
                placeholder="绘制数字后点击“开始识别”...",
                interactive=False  # 禁止用户编辑
            )
            # 额外提示
            gr.Markdown("""
            ### 📝 使用建议
            1. 用鼠标在画布中央写数字（0-9），笔画尽量粗且连贯
            2. 避免数字过小/过大，不要超出画布边缘
            3. 若提示“置信度低”，可重新绘制更清晰的数字
            """)
    
    # 绑定按钮事件
    predict_btn.click(
        fn=predict_digit,
        inputs=canvas,
        outputs=output_text
    )
    clear_btn.click(
        fn=lambda: None,  # 清空画布
        inputs=None,
        outputs=canvas,
        queue=False
    )

# ---------------------- 5. 启动应用 ----------------------
if __name__ == "__main__":
    print("🚀 手写数字识别应用启动中...")
    print("📌 访问 http://localhost:7860 即可使用")
    iface.launch(
        debug=False,  # 关闭调试模式（避免异步问题）
        server_name="localhost",
        server_port=7860,
        show_api=False  # 不显示API文档（简化用户体验）
    )